package com.peng.leetcode.heap;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * TopKFrequent
 * 给定一个非空的整数数组，返回其中出现频率前 k 高的元素。
 *
 * @author: lupeng6
 * @create: 2020/12/11 18:03
 */
public class TopKFrequent {

    public static void main(String[] args) {
        int[] arr = {1, 1, 1, 2, 2, 3};
        int k = 2;
        int[] ints = new Solution().topKFrequent(arr, k);
        Arrays.stream(ints).forEach(System.out::println);
    }

    static class Solution {
        public int[] topKFrequent(int[] nums, int k) {
            Map<Integer, Integer> map = new HashMap<>();
            for (int i = 0; i < nums.length; i++) {
                map.put(nums[i], map.getOrDefault(nums[i], 0) + 1);
            }

            PriorityQueue<int[]> priorityQueue = new PriorityQueue<>((a, b) -> b[1] - a[1]);

            for (Map.Entry<Integer, Integer> entry : map.entrySet()) {
                priorityQueue.offer(new int[]{entry.getKey(), entry.getValue()});
            }
            int[] retArr = new int[k];
            for (int i = 0; i < k; i++) {
                retArr[i] = priorityQueue.poll()[0];
            }
            return retArr;
        }
    }

}
